/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "commandline.hh"
#include "learningrate.hh"
#include "loss.hh"
#include "mode.hh"

#include "iwrestimateworker.hh"
#include "iwrparseworker.hh"
#include "iwrreportworker.hh"
#include "regressor.hh"

#include <boost/format.hpp>
#include <boost/optional.hpp>
#include <boost/none.hpp>
#include <boost/program_options/options_description.hpp>
#include <boost/program_options/parsers.hpp>
#include <boost/program_options/positional_options.hpp>
#include <boost/program_options/variables_map.hpp>
#include <exception>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>

namespace
{
using namespace flassol;

int
input_loop (std::istream&                       in,
            boost::optional<std::ostream&>      report,
            boost::optional<std::ostream&>     	predict,
            bool                                test_only,
            Regressor&                          r,
            std::auto_ptr<Loss>&                loss,
            std::auto_ptr<LearningRate>&        eta,
            float                               lambda,
            bool                                adaptive,
            unsigned int                        num_threads)
{
  enum StandardParse::hash_option_value ho = 
    static_cast<StandardParse::hash_option_value> 
      (r.get_header ().hash_option);

  std::vector<RegressionQueue*> rqvector (num_threads);
  for (unsigned int i = 0; i < num_threads; ++i)
    {
      rqvector[i] = new RegressionQueue (50);
    }
  Tee<RegressionQueue> rqtee (rqvector);

  IwrParseWorker parser (ho, in, rqtee, adaptive);
  boost::thread parser_thread (boost::ref (parser));

  boost::thread_group worker_threads;

  RegressionQueue report_queue (50);

  for (unsigned int i = 0; i < num_threads; ++i)
    {
      worker_threads.create_thread
	(IwrEstimateWorker (*(rqvector[i]),
                            report_queue,
                            r,
                            i,
                            num_threads,
                            *eta,
                            *loss,
                            lambda,
                            adaptive,
                            test_only));
    }

  IwrReportWorker reporter (report_queue, report, predict, *loss);
  boost::thread reporter_thread (boost::ref (reporter));

  parser_thread.join ();
  worker_threads.join_all ();

  for (unsigned int i = 0; i < num_threads; ++i)
    {
      delete rqvector[i];
    }

  reporter_thread.join ();

  return 0;
}

int
common_main (int argc,
             char* argv[],
             boost::program_options::options_description& desc,
             boost::program_options::options_description& all,
             const boost::program_options::positional_options_description& pd)
{
  using namespace boost::program_options;
  using std::cerr;
  using std::endl;
  using std::exception;
  using std::string;
  using std::vector;

  desc.add_options ()
    ("help", "produce help message")
    ("adaptive", value<bool> ()->default_value (true), "use adaptive learning rate")
    ("eta", value<string> ()->default_value ("powerlaw (1, 1, 0)"), "learning rate")
    ("input", value<string> (), "(optional) file to read input")
    ("lambda", value<float> ()->default_value (0.1f), "dyadic l2 regularizer")
    ("model", value<string>()->default_value ("model"), "model file")
    ("num_threads", value<unsigned int> ()->default_value (1), "number of threads (doesn't seem to speed things up at the moment)")
    ("num_weight_bits", value<unsigned int>()->default_value (16), "log_2 number of weights")
    ("predict", value<string>(), "(optional) file to output predictions")
    ("ngram", value<vector<string> > (), "ngram spec")
    ("quadratic", value<vector<string> > (), "quadratic interactions spec")
    ("dotproduct", value<vector<string> > (), "dotproduct interactions spec")
    ("quiet", "suppress output")
    ("test", "only test: do not learn")
    ;
  all.add (desc);

  variables_map vm;

  try
    { 
      command_line_parser parser (argc, argv);
      parser.options (all);
      parser.positional (pd);
      store (parser.run (), vm);
      notify (vm);
    }
  catch (exception& e)
    { 
      cerr << "ERROR: " << e.what () << endl;
      cerr << desc;
      return 1;
    }

  if (vm.count ("help"))
    { 
      if (vm["mode"].as<string> () == "iwr")
        {
          cerr << "importance weighted regression" << endl;
          cerr << "  estimate a real-valued quantity which (hopefully) minimizes " << endl;
          cerr << "  the conditional expected loss given the input" << endl << endl;
        }
      else 
        {
          cerr << "importance weighted (binary) classification" << endl;
          cerr << "  estimate a class label which (hopefully) minimizes " << endl;
          cerr << "  the conditional expected weighted loss given the input" << endl << endl;
        }
      cerr << desc;
      return 1;
    }

  try
    {
      std::auto_ptr<Loss> tmp = get_loss (vm["loss"].as<string> ());
    }
  catch (std::invalid_argument& ia)
    {
      cerr << "ERROR: invalid loss specification: '" 
           << vm["loss"].as<string> ()
           << "': " 
           << ia.what () 
           << endl;
      return 1;
    }

  if (vm["num_threads"].as<unsigned int> () == 0)
    {
      cerr << "ERROR: num_threads underflow (= "
           << vm["num_threads"].as<unsigned int> ()
           << " )" << endl;
      return 1;
    }

  if (vm["num_weight_bits"].as<unsigned int> () > 63)
    {
      cerr << "ERROR: num_weight_bits overflow (= " 
           << vm["num_weight_bits"].as<unsigned int> ()
           << " )" << endl;
      return 1;
    }

  std::vector<std::pair<uint8_t, Regressor::NGramSpec> > ngram;

  if (vm.count ("ngram"))
    {
      vector<string> ns = vm["ngram"].as<vector<string> > ();

      for (vector<string>::iterator n = ns.begin (); n != ns.end (); ++n)
        {
          boost::optional<std::pair<uint8_t, Regressor::NGramSpec> > 
            parsed_n = parse_ngram (*n);

          if (! parsed_n)
            {
              return 1;
            }

          ngram.push_back (*parsed_n);
        }
    }

  std::vector<std::pair<uint8_t, uint8_t> > quadratic;

  if (vm.count ("quadratic"))
    {
      vector<string> qs = vm["quadratic"].as<vector<string> > ();

      for (vector<string>::iterator q = qs.begin (); q != qs.end (); ++q)
        {
          boost::optional<std::pair<uint8_t, uint8_t> > 
            parsed_q = parse_quadratic (*q);

          if (! parsed_q)
            {
              return 1;
            }

          quadratic.push_back (*parsed_q);
        }
    }

  std::vector<Regressor::DotProductSpec> dotproduct;

  if (vm.count ("dotproduct"))
    {
      vector<string> qs = vm["dotproduct"].as<vector<string> > ();

      for (vector<string>::iterator q = qs.begin (); q != qs.end (); ++q)
        {
          boost::optional<Regressor::DotProductSpec> 
            parsed_q = parse_dotproduct (*q);

          if (! parsed_q)
            {
              return 1;
            }

          dotproduct.push_back (*parsed_q);
        }
    }

  std::auto_ptr<LearningRate> eta;
  
  try
    {
      std::auto_ptr<LearningRate> tmp = get_learning_rate (vm["eta"].as<string> ());
      eta = tmp;
    }
  catch (std::invalid_argument& ia)
    {
      cerr << "ERROR: invalid learning rate specification '" 
           << vm["eta"].as<string> ()
           << "': " 
           << ia.what () 
           << endl;
      return 1;
    }

  if (vm["lambda"].as<float> () < 0.0f)
    {
      cerr << "ERROR: invalid dyadic regularizer specification "
           << vm["lambda"].as<float> ()
           << endl;
      return 1;
    }

  Regressor r = Regressor::Open (vm["model"].as<string> ())
                  .num_weights (1 << vm["num_weight_bits"].as<unsigned int> ())
                  .loss (vm["loss"].as<string> ())
                  .read_only (vm.count ("test"))
                  .add_quadratic (quadratic)
                  .add_dotproduct (dotproduct)
                  .add_ngram (ngram)
                ;

  if (r.get_header ().num_classes != 2)
    {
      cerr << "ERROR: multiclass model detected: num_classes = " 
           << r.get_header ().num_classes
           << std::endl;

      return 1;
    }


  std::auto_ptr<Loss> loss;
  
  try
    {
      std::auto_ptr<Loss> tmp = get_loss (r.get_header ().loss);
      loss = tmp;
    }
  catch (std::invalid_argument& ia)
    {
      cerr << "ERROR: invalid loss specification: '" 
           << r.get_header ().loss 
           << "': " 
           << ia.what () 
           << endl;
      return 1;
    }

  std::ifstream input;

  if (vm.count ("input"))
    {
      input.open (vm["input"].as<string> ().c_str ());

      if (! input.good ())
        {
          cerr << "ERROR: can't open input file " 
               << vm["input"].as<string> ()
               << ": " << strerror (errno)
               << endl;
    
          return 1;
        }
    }

  std::ofstream predict;

  if (vm.count ("predict"))
    {
      predict.open (vm["predict"].as<string> ().c_str ());

      if (! predict.good ())
        {
          cerr << "ERROR: can't open prediction file " 
               << vm["predict"].as<string> ()
               << ": " << strerror (errno)
               << endl;

          return 1;
        }
    }

  if (vm.count ("quiet") == 0)
    {
      cerr << "adaptive = " << (vm["adaptive"].as<bool> () ? 1 : 0) << endl;
      cerr << "eta = " << vm["eta"].as<string> () << endl;
      cerr << "lambda = " << vm["lambda"].as<float> () << endl;
      cerr << "num_threads = " << vm["num_threads"].as<unsigned int> () << endl;
      if (vm.count ("test"))
        {
          cerr << "test only" << endl;
        }

      if (vm.count ("input"))
        {
          cerr << "read input from " 
               << vm["input"].as<string> ()
               << endl;
        }

      if (vm.count ("predict"))
        {
          cerr << "write predictions to " 
               << vm["predict"].as<string> ()
               << endl;
        }

      cerr << r.get_header ();
    }

  return input_loop ((vm.count ("input") 
                        ? static_cast<std::istream&> (input)
                        : std::cin),
                     (vm.count ("quiet") 
                        ? boost::optional<std::ostream&> (boost::none)
                        : boost::optional<std::ostream&> (cerr)),
                     (vm.count ("predict")
                        ? boost::optional<std::ostream&> (static_cast<std::ostream&> (predict))
                        : boost::optional<std::ostream&> (boost::none)),
                     vm.count ("test"),
                     r, 
                     loss,
                     eta,
                     vm["lambda"].as<float> (),
                     vm["adaptive"].as<bool> (),
                     vm["num_threads"].as<unsigned int> ());
}

}

namespace flassol
{

int
iwr_main (int   argc,
          char* argv[])
  {
    using namespace boost::program_options;
    using std::string;
    using std::vector;

    options_description desc ("Allowed options");
    options_description all ("");
    positional_options_description pd;

    desc.add_options ()
      ("loss", value<string>()->default_value ("squaredclipped"), get_registered_losses ().c_str ());
    all.add_options ()
      ("mode", value<string> ()->default_value ("iwr"), modes_string.c_str ());

    pd.add ("mode", 1);
    pd.add ("model", -1);

    return common_main (argc, argv, desc, all, pd);
  }

int
iwc_main (int   argc,
          char* argv[])
  {
    using namespace boost::program_options;
    using std::string;
    using std::vector;

    options_description desc ("Allowed options");
    options_description all ("");
    positional_options_description pd;

    desc.add_options ()
      ("loss", value<string> ()->default_value ("hinge"), get_registered_losses ().c_str ());
    all.add_options ()
      ("mode", value<string> ()->default_value ("iwc"), modes_string.c_str ());

    pd.add ("mode", 1);
    pd.add ("model", -1);

    return common_main (argc, argv, desc, all, pd);
  }

}
